import tushare as ts
import pandas as pd
import backtrader as bt


# 1. 数据准备阶段

# 配置Tushare token
ts.set_token('你的Tushare Token')
pro = ts.pro_api()

# 获取日线数据（示例：000001.SZ 2025年数据）
df = pro.daily(ts_code='000001.SZ', start_date='20250101', end_date='20250228')
df = df.sort_values('trade_date').reset_index(drop=True)

# 转换为Backtrader所需格式
df['trade_date'] = pd.to_datetime(df['trade_date'])
df.rename(columns={
    'trade_date': 'date',
    'open': 'open',
    'high': 'high',
    'low': 'low',
    'close': 'close',
    'vol': 'volume'
}, inplace=True)
df = df[['date', 'open', 'high', 'low', 'close', 'volume']]
df.set_index('date', inplace=True)

# 保存到CSV（供Backtrader读取）
df.to_csv('000001.SZ.csv')


# 2. 信号生成模块

# 假设每日分析结果保存为如下格式的CSV
signals = pd.DataFrame({
    'date': ['2025-02-26', '2025-02-27'],
    'operation_advice': ['轻仓策略', '逢高减仓'],
    'confidence_level': ['B', 'C'],
    'prediction': ['震荡偏弱', '下跌']
})

# 映射操作建议到仓位比例
signal_mapping = {
    '轻仓策略': 0.3,
    '逢高减仓': 0.1,
    '全仓': 1.0
}
signals['position_ratio'] = signals['operation_advice'].map(signal_mapping)

# 合并行情数据与信号数据
merged_data = pd.merge(df, signals, left_index=True, right_on='date', how='left')
merged_data.to_csv('merged_data.csv')


# 3. Backtrader策略实现
class SignalStrategy(bt.Strategy):
    params = (
        ('position_ratio', 0.3),  # 默认仓位比例
    )

    def __init__(self):
        self.signal = self.datas[0].signal  # 信号数据列
        self.order = None

    def next(self):
        # 检查是否有信号
        if self.signal[0] == 0:
            return

        # 清空当前持仓
        if self.order:
            return
        if self.position:
            self.close()

        # 执行开仓
        target_value = self.broker.getvalue() * self.signal[0]
        self.order = self.order_target_value(target=target_value)

# 添加信号数据到K线
class SignalData(bt.feeds.PandasData):
    lines = ('signal',)  # 新增信号线
    params = (
        ('datetime', None),
        ('open', 0),
        ('high', 1),
        ('low', 2),
        ('close', 3),
        ('volume', 4),
        ('signal', 5),
        ('openinterest', -1)
    )


# 4. 回测执行代码
def run_backtest():
    cerebro = bt.Cerebro()

    # 加载数据
    data = SignalData(
        dataname=merged_data,
        fromdate=pd.to_datetime('2025-01-01'),
        todate=pd.to_datetime('2025-02-28')
    )
    cerebro.adddata(data)

    # 添加策略
    cerebro.addstrategy(SignalStrategy)

    # 设置初始资金和佣金
    cerebro.broker.setcash(1000000)
    cerebro.broker.setcommission(commission=0.001)  # 0.1%佣金

    # 添加分析器
    cerebro.addanalyzer(bt.analyzers.SharpeRatio, _name='sharpe')
    cerebro.addanalyzer(bt.analyzers.DrawDown, _name='drawdown')

    # 运行回测
    results = cerebro.run()
    strat = results[0]

    # 打印结果
    print(f'期末总资产: {cerebro.broker.getvalue():.2f}')
    print(f'夏普比率: {strat.analyzers.sharpe.get_analysis()["sharperatio"]:.2f}')
    print(f'最大回撤: {strat.analyzers.drawdown.get_analysis()["max"]["drawdown"]:.2f}%')

    # 可视化
    cerebro.plot(style='candlestick', volume=False)


if __name__ == '__main__':
    # 步骤1：获取并预处理数据
    df = get_tushare_data()
    signals = load_analysis_signals()
    merged_data = merge_data(df, signals)

    # 步骤2：运行回测
    run_backtest()
